import torch

class Config:
    # 数据集配置
    DATASET = "CIFAR10"
    DATA_DIR = "./data"
    NUM_CLASSES = 10
    
    # 训练配置
    BATCH_SIZE = 128
    EPOCHS = 100
    LEARNING_RATE = 0.1
    MOMENTUM = 0.9
    WEIGHT_DECAY = 5e-4
    LR_SCHEDULER = "cosine"  # "step" or "cosine"
    
    # 模型配置
    MODEL_NAME = "ResNet18"
    PRETRAINED = False
    
    # 设备配置
    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
    NUM_WORKERS = 4
    
    # 日志和保存
    LOG_DIR = "./logs"
    SAVE_DIR = "./checkpoints"
    SAVE_FREQ = 10  # 每多少epoch保存一次模型
    
    # 数据增强
    USE_CUTMIX = True
    USE_MIXUP = True
    USE_AUTOAUGMENT = True

config = Config()